import numpy as np
import matplotlib.pyplot as plt

# 使用K-means底层，完成以下操作
# 数据集[[1, 2], [2, 2], [6, 8], [7, 8]]
# 1.	将数据集使用合适方式进行处理（10分）
X = np.array([[1, 2], [2, 2], [6, 8], [7, 8]], dtype=np.float64)
# from sklearn.preprocessing import StandardScaler
# X = StandardScaler().fit_transform(X)
m, n = X.shape
k = 2

# 2.	使用随机方式，初始化质心（10分）
idx_rnd = np.random.randint(0, m, k)
C = X[idx_rnd]
print(C)

# 3.	完成模型计算
# a)	设置循环5次（10分）
iters = 0
while (iters < 5):
    iters += 1

    # b)	计算数据集中每个样本点到各个质心的距离，添加到列表中（20分）
    dis_mat = np.zeros([m, k])
    for i in range(k):
        dis_mat[:, i] = ((X - C[i])**2).sum(axis=1)**0.5

    # c)	计算距离距离哪个样本最近（20分）
    class_vector = dis_mat.argmin(axis=1)

    # d)	使用均值更新质心（10分）
    for i in range(k):
        C[i] = X[class_vector == i].mean(axis=0)

# 4.	打印样本的归属于哪一个簇（10分）
cmap = plt.get_cmap('rainbow', k)
for i in range(k):
    idx = class_vector == i
    plt.scatter(X[idx, 0], X[idx, 1], color=cmap(i), label=f'class#{i + 1}')

# 5.	打印质心位置（10分）
for i in range(k):
    plt.scatter(C[i, 0], C[i, 1], color=cmap(i), s=100, zorder=100, marker='x')
    plt.annotate(f'center#{i + 1}', xy=[C[i, 0], C[i, 1]], color=cmap(i))

# finally show all plotting
plt.legend()
plt.show()
